import cv2
import numpy as np


def SNR(img: cv2.Mat, img_noise: cv2.Mat):
    """
    计算信噪比
    """
    img = img.astype(np.float32)
    img_noise = img_noise.astype(np.float32)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_noise = cv2.cvtColor(img_noise, cv2.COLOR_BGR2GRAY)
    img = img.flatten()
    img_noise = img_noise.flatten()
    img_var = np.var(img)
    img_noise_var = np.var(img_noise)
    SNR = 20 * np.log10(np.sqrt(img_var) / np.sqrt(img_noise_var))
    return SNR


if __name__ == "__main__":
    img = cv2.imread("../bg_img/background.png")
    img1 = img.copy()
    # add gaussian blur
    cv2.GaussianBlur(img1, (5, 5), 2, dst=img1)
    # cv2.imwrite("../bg_img/background_blur.png", img1)
    img2 = img1.copy()
    # add median blur
    img2 = np.random.randn(img.shape[0], img.shape[1], img.shape[2]) * 50 + img2 - 20
    snr = SNR(img, img2)
    print("SNR:{0}", snr)
    cv2.imwrite(f"../bg_img/background_snr_{snr:.3f}.png", img2)
